# -*- coding: utf-8 -*-
# @Time    : 2021/10/28 16:47
# @Author  : huangwei
# @File    : img2excel2.py
# @Software: PyCharm
import cv2
import time
from config_args import args
from excel_method import split_floors, get_table_box_data, write_excel
from text import TextDetector, TextRecognizer
from method import get_lines, get_text_boxes, roll_image_horize, sort_line, alter_close_lines, fix_up_lines, line_cross, \
    alter_lines, fill_rect_table
from draw_method import draw_text_boxes, draw_table, draw_img
from table_line_net import table_net

# 动态分配内存
import tensorflow as tf

config = tf.compat.v1.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.5
config.gpu_options.allow_growth = True
session = tf.compat.v1.InteractiveSession(config=config)

# 检测文本框和识别文字方法
text_detector = TextDetector(args)
text_recognizer = TextRecognizer(args)

# 加载识别表格线的模型
table_line_model_path = 'det_table_line_models/table-line.h5'
table_line_model = table_net((None, None, 3), 2)
table_line_model.load_weights(table_line_model_path)

if __name__ == '__main__':
    file_path = "images/261.png"
    # file_path = "images/label123.png"
    roll_path = "images/roll_img.png"

    start_time = time.time()

    roll_image_horize(file_path, roll_path, table_line_model)
    img = cv2.imread(roll_path)
    h, w = img.shape[:2]

    # 找到所有的横线和竖线，去除掉不平行的，留下的调整为水平或者竖直的
    input_size = (1024, 1024)
    row_lines, col_lines = get_lines(roll_path, table_line_model, input_size)

    # 调整为水平或者竖直
    row_lines = alter_lines(row_lines)
    col_lines = alter_lines(col_lines, axis=1)
    draw_img((h, w), row_lines + col_lines, 'temp_path/dd.png')

    # 两条竖线或两条横线贴近的去除一条
    for i in range(2):
        row_lines = alter_close_lines(row_lines)
        col_lines = alter_close_lines(col_lines, axis=1)

    # 竖线直线之间是否相交及相交的点往外突出一些
    row_lines, col_lines = line_cross(row_lines, col_lines)

    # 对线进行从左到右从上到下排序
    row_lines = sort_line(row_lines)
    col_lines = sort_line(col_lines, axis=1)
    draw_img((h, w), row_lines + col_lines, 'temp_path/ee.png')
    print("排序完用时:", time.time() - start_time)

    # 识别文本框
    det_text_boxes = get_text_boxes(img, text_detector, row_lines)
    print("文本框的数量为：", len(det_text_boxes))
    text_box_img_path = "temp_path/text_box.png"
    draw_text_boxes(img, det_text_boxes, save_path=text_box_img_path)
    print("识别完文字框用时:", time.time() - start_time)

    # 依次扫描每一条横线和竖线，判断其是否需要去除冒头的部分或延长冒头的部分。
    print("调整前横线竖线数量:", len(row_lines), len(col_lines))
    row_lines, col_lines = fix_up_lines(row_lines, col_lines, det_text_boxes, roll_path)
    draw_img((h, w), row_lines + col_lines, 'temp_path/ff.png')
    print("调整线用时:", time.time() - start_time)

    # 补全最外围的线,即确保为矩形
    row_lines, col_lines = fill_rect_table(row_lines, col_lines)
    draw_img((h, w), row_lines + col_lines, 'temp_path/gg.png')

    # 两条竖线或两条横线贴近的去除一条
    row_lines = alter_close_lines(row_lines)
    col_lines = alter_close_lines(col_lines, axis=1)
    row_lines = sort_line(row_lines)
    col_lines = sort_line(col_lines, axis=1)
    draw_img((h, w), row_lines + col_lines, 'temp_path/hh.png')
    print("调整后横线竖线数量:", len(row_lines), len(col_lines))
    print("调整后横线用时:", time.time() - start_time)

    # 画出识别出的线识别出的表格
    table_boxes = draw_table(img, row_lines + col_lines, 'temp_path/table_box.png')

    # 将数据存储到excel中
    # 表格分层
    x_index, y_index, row_lines, col_lines = split_floors(row_lines, col_lines)

    # 获取表中提取的信息
    box_infos = get_table_box_data(roll_path, text_recognizer, table_boxes, det_text_boxes, x_index, y_index)

    write_excel(box_infos, x_index, y_index, 'temp_path/merging.xlsx')
    print("完成用时:", time.time() - start_time)
